import pymysql
from neo4j import GraphDatabase
from openai import OpenAI  # 使用 OpenAI 客户端
import time
'''
此脚本连接MySQL数据库，
获取表结构信息，
调用OpenAI兼容接口生成表的嵌入向量，
并将这些信息导入到Neo4j数据库中。
'''
# ==================== 配置 ====================
DB_CFG = dict(
    host='127.0.0.1',
    port=3306,
    user='root',
    password='12345678',
    charset='utf8mb4'
)

NEO4J_URI = "bolt://localhost:7687"
NEO4J_AUTH = ("neo4j", "12345678")

# 🔁 OpenAI 兼容 Embedding 接口配置
EMBEDDING_BASE_URL = "https://api.siliconflow.cn/v1"  # 修改为你的实际地址
EMBEDDING_MODEL = "BAAI/bge-m3"                 # 或 text-embedding-3-small
EMBEDDING_API_KEY = "sk-xpczcmiojslifyfbfzkrkctbgbilljxoekpevjhcdnrcdzjb"                     # 多数本地服务设为 "EMPTY"

# 初始化 OpenAI 客户端
client = OpenAI(api_key=EMBEDDING_API_KEY, base_url=EMBEDDING_BASE_URL)
# =============================================

def get_table_semantic_text(table_info, columns):
    """
    构造表的语义文本，用于 embedding
    示例："users 用户表 包含字段: id, name, email"
    """
    name = table_info['TABLE_NAME']
    comment = table_info.get('TABLE_COMMENT', '') or ""
    col_names = ", ".join([col['COLUMN_NAME'] for col in columns[:5]])
    if len(columns) > 5:
        col_names += " ..."
    return f"{name} {comment} 包含字段: {col_names}".strip()


def get_embeddings(texts, batch_size=100):
    """
    调用 OpenAI 兼容接口批量生成 embedding
    """
    embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        print(f"🧠 调用 embedding API ({i + 1} ~ {i + len(batch)})...")

        try:
            response = client.embeddings.create(
                input=batch,
                model=EMBEDDING_MODEL
            )
            # 提取 embedding（注意：response.data 是 list）
            batch_embeddings = [d.embedding for d in response.data]
            embeddings.extend(batch_embeddings)
        except Exception as e:
            print(f"❌ 调用 embedding API 失败: {e}")
            # 可以选择跳过或重试
            embeddings.extend([None] * len(batch))  # 占位

        time.sleep(0.1)  # 避免请求过快

    return embeddings


def init_all_tables():
    print("🔄 正在连接 MySQL...")
    try:
        conn = pymysql.connect(**DB_CFG)
        cur = conn.cursor(pymysql.cursors.DictCursor)
        print("✅ MySQL 连接成功！")
    except Exception as e:
        print(f"❌ 无法连接 MySQL：{e}")
        return

    INCLUDED = ('test',)
    print(f"🔍 正在查询不属于 {INCLUDED} 的业务表...")

    try:
        cur.execute("""
            SELECT TABLE_SCHEMA, TABLE_NAME, TABLE_COMMENT
            FROM information_schema.TABLES
            WHERE TABLE_SCHEMA IN %s AND TABLE_TYPE = 'BASE TABLE'
        """, (INCLUDED,))
        tables = cur.fetchall()
    except Exception as e:
        print(f"❌ 查询失败：{e}")
        cur.close()
        conn.close()
        return

    if not tables:
        print("🚫 无可用表，请检查数据库内容。")
        cur.close()
        conn.close()
        return

    print(f"📊 发现 {len(tables)} 个表，正在构造语义文本...")

    # Step 1: 构造每张表的语义文本
    all_texts = []
    table_data = []

    for t in tables:
        schema, tbl = t['TABLE_SCHEMA'], t['TABLE_NAME']
        cur.execute("""
            SELECT COLUMN_NAME, DATA_TYPE, COLUMN_COMMENT
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
        """, (schema, tbl))
        columns = cur.fetchall()

        semantic_text = get_table_semantic_text(t, columns)
        all_texts.append(semantic_text)
        table_data.append((schema, tbl, t['TABLE_COMMENT'], columns))

    # Step 2: 调用 OpenAI 兼容接口获取 embedding
    print("🔄 正在调用 embedding 接口...")
    embeddings = get_embeddings(all_texts)

    if len([e for e in embeddings if e is None]) == len(embeddings):
        print("❌ 所有 embedding 生成失败，终止导入。")
        return

    print("✅ embedding 获取完成，开始导入 Neo4j...")

    # Step 3: 写入 Neo4j
    with GraphDatabase.driver(NEO4J_URI, auth=NEO4J_AUTH).session() as s:
        try:
            for i, (schema, tbl, comment, columns) in enumerate(table_data):
                embedding = embeddings[i]
                if embedding is None:
                    print(f"⚠️  embedding 失败，跳过表: {schema}.{tbl}")
                    continue

                # 创建 Database
                s.run("MERGE (d:Database {name: $db})", db=schema)

                # 创建 Table（带 embedding）
                s.run("""
                    MERGE (t:Table {schema: $db, name: $tbl})
                    SET t.comment = $tbl_comment,
                        t.embedding = $embedding
                    WITH t
                    MATCH (d:Database {name: $db})
                    MERGE (d)-[:HAS_TABLE]->(t)
                """, db=schema, tbl=tbl, tbl_comment=comment or "", embedding=embedding)

                # 创建字段节点
                cur.execute("""
                    SELECT COLUMN_NAME
                    FROM information_schema.KEY_COLUMN_USAGE
                    WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s AND CONSTRAINT_NAME = 'PRIMARY'
                """, (schema, tbl))
                primary_keys = {row['COLUMN_NAME'] for row in cur.fetchall()}

                for col in columns:
                    col_name = col['COLUMN_NAME']
                    is_pk = col_name in primary_keys
                    s.run("""
                        MERGE (c:Column {schema: $db, table: $tbl, name: $col})
                        SET c.type = $typ,
                            c.comment = $col_comment,
                            c.is_primary = $is_pk
                        WITH c, $is_pk AS pk_flag
                        MATCH (t:Table {schema: $db, name: $tbl})
                        MERGE (t)-[:HAS_COLUMN]->(c)
                        FOREACH (_ IN CASE WHEN pk_flag THEN [1] ELSE [] END |
                            MERGE (t)-[:HAS_PRIMARY_KEY]->(c)
                        )
                    """,
                        db=schema,
                        tbl=tbl,
                        col=col_name,
                        typ=col['DATA_TYPE'],
                        col_comment=col['COLUMN_COMMENT'] or "",
                        is_pk=is_pk
                    )

                if (i + 1) % 10 == 0 or i == len(tables) - 1:
                    print(f"📦 已处理 {i + 1}/{len(tables)} 个表...")

            # Step 4: 创建向量索引
            print("🏗️ 正在创建向量索引...")
            try:
                dim = len(embeddings[0])
                s.run("""
                CREATE VECTOR INDEX table_comment_embedding 
                FOR (t:Table) ON (t.embedding) 
                OPTIONS { indexConfig: {
                  `vector.dimensions`: $dims,
                  `vector.similarity_function`: 'cosine'
                }}
                """, dims=dim)
                print("✅ 向量索引创建成功！")
            except Exception as e:
                if "An equivalent index already exists" in str(e):
                    print("🟡 向量索引已存在，跳过。")
                else:
                    print(f"⚠️ 创建索引失败：{e}")

            print("🎉 所有表结构与 embedding 已导入 Neo4j！")
            print("✅ 可用于后续：语义搜索、LLM 自动连接 Answer → Table")

        except Exception as e:
            print(f"❌ 写入 Neo4j 失败：{e}")
            import traceback
            traceback.print_exc()

    cur.close()
    conn.close()


# ============ 主程序入口 ============
if __name__ == '__main__':
    try:
        init_all_tables()
    except KeyboardInterrupt:
        print("\n👋 用户中断。")
    except Exception as e:
        print(f"❌ 异常：{e}")
        import traceback
        traceback.print_exc()